#include "ntifs.h"
#include "ntddk.h"
#include "ia32_types.h"
#include "ept.h"
#include "AsmCall.h"
#include "common.h"
#include "Util.h"
#include "Vmm.h"
#include "HooksFunction.h"

PHOOK_EPT_STRUCT GetFunctionHookedStruct(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pPa) {
	if (pStruct == 0 || *pStruct == 0)
	{
		return 0;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->pOriginalPA == pPa)
		{
			return i;
		}
	}

	return 0;
}

bool InsertHookStructIntoList(PHOOK_EPT_STRUCT* pStruct, PHOOK_EPT_STRUCT pInsertStruct) {
	if (pStruct == 0 || pInsertStruct == 0)
	{
		return false;
	}

	pInsertStruct->pLast = 0;
	if (*pStruct != 0)
	{
		(*pStruct)->pLast = pInsertStruct;
	}
	pInsertStruct->pNext = *pStruct;
	*pStruct = pInsertStruct;

	return true;
}

bool FindFunctionIsHooked(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pPa) {

	if (pStruct == 0 || *pStruct == 0)
	{
		return false;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->pOriginalPA == pPa)
		{
			return true;
		}
	}

	return false;
}

ULONG_PTR GetRing3HookByPIDAndOriginalVA(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pOriginalVa, ULONG_PTR nProcessID) {
	if (pStruct == 0 || *pStruct == 0)
	{
		return 0;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->bUsingNow)
		{
			if (i->nProcessPid == nProcessID)
			{
				if (i->pOriginalVA == pOriginalVa)
				{
					return i->pHookVA;
				}
			}
		}
	}

	return 0;
}

bool FindRing3FunctionIsHooked(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pOriginalVa, ULONG_PTR nProcessID) {

	if (pStruct == 0 || *pStruct == 0)
	{
		return false;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->bUsingNow)
		{
			if (i->nProcessPid == nProcessID)
			{
				if (i->pOriginalVA == pOriginalVa)
				{
					return true;
				}
			}
		}
	}

	return false;
}

bool FindRing3UserProcIsHooked(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR nProcessID) {

	if (pStruct == 0 || *pStruct == 0)
	{
		return false;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->bUsingNow)
		{
			if (i->nProcessPid == nProcessID)
			{
				return true;
			}
		}
	}

	return false;
}

PHOOK_EPT_STRUCT RemoveHookStructByPA(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pa) {
	PHOOK_EPT_STRUCT pReturn = 0;
	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->pOriginalPA == pa)
		{
			if (i->pLast)
			{
				((PHOOK_EPT_STRUCT)(i->pLast))->pNext = i->pNext;
			}
			else {
				*pStruct = (PHOOK_EPT_STRUCT)(i->pNext);
			}

			if (i->pNext)
			{
				((PHOOK_EPT_STRUCT)(i->pNext))->pLast = i->pLast;
			}

			pReturn = i;
			break;
		}
	}

	return pReturn;
}

bool FullEptStructAccess(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pOriginalPa, EptCommonEntry *pEptCommonEntry) {

	if (pStruct == 0 || *pStruct == 0 || pEptCommonEntry == 0)
	{
		if (pEptCommonEntry != 0)
		{
			pEptCommonEntry->fields.read_access = 1;
			UtilInveptGlobal();
		}
		return false;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->bUsingNow)
		{
			if (i->pOriginalPA == pOriginalPa)
			{
				if (i->pOriginalPA == pEptCommonEntry->fields.physial_address)
				{
					pEptCommonEntry->fields.physial_address = i->pHookPA;
					pEptCommonEntry->fields.read_access = i->Hook_read_access;
					pEptCommonEntry->fields.execute_access = i->Hook_execute_access;
					pEptCommonEntry->fields.write_access = i->Hook_write_access;
					pEptCommonEntry->fields.memory_type = i->Hook_memory_type;
					return true;
				}
				else if (i->pHookPA == pEptCommonEntry->fields.physial_address) {
					pEptCommonEntry->fields.physial_address = i->pOriginalPA;
					pEptCommonEntry->fields.read_access = i->Original_read_access;
					pEptCommonEntry->fields.execute_access = i->Original_execute_access;
					pEptCommonEntry->fields.write_access = i->Original_write_access;
					pEptCommonEntry->fields.memory_type = i->Original_memory_type;
					return true;
				}
			}
		}
	}

	pEptCommonEntry->fields.read_access = 1;
	UtilInveptGlobal();
	return false;
}

bool FullEptStructWriteAccess(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pOriginalPa, EptCommonEntry* pEptCommonEntry, ULONG_PTR pTotalOriginalPa) {

	if (pStruct == 0 || *pStruct == 0 || pEptCommonEntry == 0)
	{
		if (pEptCommonEntry != 0)
		{
			pEptCommonEntry->fields.write_access = 1;
			UtilInveptGlobal();
		}
		return false;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->bUsingNow)
		{
			if (i->pOriginalPA == pOriginalPa)
			{
				if (i->pOriginalPA == pEptCommonEntry->fields.physial_address)
				{
					pEptCommonEntry->fields.physial_address = i->pHookPA;
					pEptCommonEntry->fields.read_access = i->Hook_read_access;
					pEptCommonEntry->fields.execute_access = i->Hook_execute_access;
					pEptCommonEntry->fields.write_access = i->Hook_write_access;
					pEptCommonEntry->fields.memory_type = i->Hook_memory_type;
					return true;
				}
				else if (i->pHookPA == pEptCommonEntry->fields.physial_address) {
					pEptCommonEntry->fields.physial_address = i->pOriginalPA;
					pEptCommonEntry->fields.read_access = i->Original_read_access;
					pEptCommonEntry->fields.execute_access = i->Original_execute_access;
					pEptCommonEntry->fields.write_access = i->Original_write_access;
					pEptCommonEntry->fields.memory_type = i->Original_memory_type;
					return true;
				}
			}
		}
	}

	pEptCommonEntry->fields.write_access = 1;
	return false;
}

bool FullExecuteEptStructAccess(PHOOK_EPT_STRUCT* pStruct, ULONG_PTR pOriginalPa, EptCommonEntry* pEptCommonEntry, ULONG_PTR pTotalOriginalPa) {

	if (pStruct == 0 || *pStruct == 0 || pEptCommonEntry == 0)
	{
		if (pEptCommonEntry != 0)
		{
			pEptCommonEntry->fields.execute_access = 1;
			UtilInveptGlobal();
		}
		return false;
	}

	for (PHOOK_EPT_STRUCT i = *pStruct; i != 0; i = (PHOOK_EPT_STRUCT)(i->pNext))
	{
		if (i->bUsingNow)
		{
			if (i->pOriginalPA == pOriginalPa)
			{
				if (i->pOriginalPA == pEptCommonEntry->fields.physial_address)
				{
					pEptCommonEntry->fields.physial_address = i->pHookPA;
					pEptCommonEntry->fields.read_access = i->Hook_read_access;
					pEptCommonEntry->fields.execute_access = i->Hook_execute_access;
					pEptCommonEntry->fields.write_access = i->Hook_write_access;
					pEptCommonEntry->fields.memory_type = i->Hook_memory_type;
					return true;
				}
				else if (i->pHookPA == pEptCommonEntry->fields.physial_address) {
					pEptCommonEntry->fields.physial_address = i->pOriginalPA;
					pEptCommonEntry->fields.read_access = i->Original_read_access;
					pEptCommonEntry->fields.execute_access = i->Original_execute_access;
					pEptCommonEntry->fields.write_access = i->Original_write_access;
					pEptCommonEntry->fields.memory_type = i->Original_memory_type;
					return true;
				}
			}
		}
	}

	pEptCommonEntry->fields.execute_access = 1;
	UtilInveptGlobal();
	return true;
}

bool HookProcFunction(ULONG_PTR pFunction, ULONG_PTR pOriginalFunction, ULONG_PTR pDealFunction, ULONG_PTR nCodeLen,ULONG_PTR pParamPointer) {
	bool bReturn = false;
	ULONG_PTR nJmpBackFunctionLen = Asm_GetJmpCodeLen();
	ULONG_PTR pJmpBackFunctionStart = Asm_GetJmpCodeStart();
	unsigned char *pJmpBackFunctionMemory = (unsigned char *)ExAllocatePoolWithTag(NonPagedPool, (nCodeLen + nJmpBackFunctionLen + 0x1000) & (~0xFFFu), kHyperPlatformCommonPoolTag);
	if (pJmpBackFunctionMemory != NULL)
	{
		RtlCopyMemory(pJmpBackFunctionMemory, (void *)pFunction, nCodeLen);
		ULONG_PTR nCodeTempLen = nCodeLen;
		unsigned char* pTempCodeBuffer = pJmpBackFunctionMemory;
		while (nCodeTempLen > 0)
		{
			ULONG_PTR nTempLen = LDE(pTempCodeBuffer);
			if (nTempLen <= nCodeTempLen)
			{
				nCodeTempLen -= nTempLen;
				if (nTempLen == 6 && pTempCodeBuffer[0] == 0x0f && pTempCodeBuffer[1] == 0x85)
				{
					ULONG_PTR nRealAddress = pTempCodeBuffer - pJmpBackFunctionMemory + pOriginalFunction;
					ULONG_PTR nOffset = *(ULONG*)&(pTempCodeBuffer[2]);
					KdPrint(("[Whisper]nRealAddress: %p %p %p %I64X %I64X \n", pTempCodeBuffer, nRealAddress, pJmpBackFunctionMemory, pOriginalFunction, nOffset));
					if (nOffset >= 0x80000000)
					{
						nRealAddress = nRealAddress + 6 - (0x100000000u - nOffset);
					}
					else {
						nRealAddress += (nOffset + 6);
					}

					unsigned char* pJmpRealAddressMemory = (unsigned char*)ExAllocatePoolWithTag(NonPagedPool, (nJmpBackFunctionLen + 0x1000) & (~0xFFFu), kHyperPlatformCommonPoolTag);
					if (pJmpRealAddressMemory != 0)
					{
						KdPrint(("[Whisper] JmpRealAddress: %p --> %p %I64X\n", pTempCodeBuffer, pJmpRealAddressMemory, nRealAddress));
						RtlCopyMemory(pJmpRealAddressMemory, (void*)pJmpBackFunctionStart, nJmpBackFunctionLen);
						*(ULONG_PTR*)(pJmpRealAddressMemory + 0x3) = nRealAddress;
						KdPrint(("[Whisper] JmpRealAddress Offset: %I64X \n", pJmpRealAddressMemory - (pTempCodeBuffer + 6)));
						*(ULONG*)&(pTempCodeBuffer[2]) = (ULONG)(pJmpRealAddressMemory - (pTempCodeBuffer + 6));
					}
				}
				else if (nTempLen == 6 && pTempCodeBuffer[0] == 0x8b && pTempCodeBuffer[1] == 0x05) {
					ULONG_PTR nRealAddress = pTempCodeBuffer - pJmpBackFunctionMemory + pOriginalFunction;
					ULONG_PTR nOffset = *(ULONG*)&(pTempCodeBuffer[2]);
					KdPrint(("[Whisper]Read dword nRealAddress: %p %p %p %I64X %I64X \n", pTempCodeBuffer, nRealAddress, pJmpBackFunctionMemory, pOriginalFunction, nOffset));
					if (nOffset >= 0x80000000)
					{
						nRealAddress = nRealAddress + 6 - (0x100000000u - nOffset);
					}
					else {
						nRealAddress += (nOffset + 6);
					}
					unsigned char* pJmpRealAddressMemory = (unsigned char*)ExAllocatePoolWithTag(NonPagedPool, (Asm_GetAddressLenAsmFunction() + 0x1000) & (~0xFFFu), kHyperPlatformCommonPoolTag);
					if (pJmpRealAddressMemory != 0)
					{
						KdPrint(("[Whisper] JmpRealAddress: %p --> %p %I64X\n", pTempCodeBuffer, pJmpRealAddressMemory, nRealAddress));
						RtlCopyMemory(pJmpRealAddressMemory, (void*)Asm_GetAddressValueAsmFunctionStart(), Asm_GetAddressLenAsmFunction());
						*(ULONG_PTR*)(pJmpRealAddressMemory + 0x2) = nRealAddress;
						KdPrint(("[Whisper] JmpRealAddress Offset: %I64X \n", pJmpRealAddressMemory - (pTempCodeBuffer + 5)));
						pTempCodeBuffer[0] = 0xE8;
						*(ULONG*)&(pTempCodeBuffer[1]) = (ULONG)(pJmpRealAddressMemory - (pTempCodeBuffer + 5));
						pTempCodeBuffer[5] = 0x90;
					}
				}
				else if (nTempLen == 7 && pTempCodeBuffer[0] == 0x48 && pTempCodeBuffer[1] == 0x8b && pTempCodeBuffer[2] == 0x05) {
					ULONG_PTR nRealAddress = pTempCodeBuffer - pJmpBackFunctionMemory + pOriginalFunction;
					ULONG_PTR nOffset = *(ULONG*)&(pTempCodeBuffer[3]);
					KdPrint(("[Whisper]Read qword nRealAddress: %p %p %p %I64X %I64X \n", pTempCodeBuffer, nRealAddress, pJmpBackFunctionMemory, pOriginalFunction, nOffset));
					if (nOffset >= 0x80000000)
					{
						nRealAddress = nRealAddress + 7 - (0x100000000u - nOffset);
					}
					else {
						nRealAddress += (nOffset + 7);
					}
					unsigned char* pJmpRealAddressMemory = (unsigned char*)ExAllocatePoolWithTag(NonPagedPool, (Asm_GetAddressQwordValueLenAsmFunction() + 0x1000) & (~0xFFFu), kHyperPlatformCommonPoolTag);
					if (pJmpRealAddressMemory != 0)
					{
						KdPrint(("[Whisper] JmpRealAddress: %p --> %p %I64X\n", pTempCodeBuffer, pJmpRealAddressMemory, nRealAddress));
						RtlCopyMemory(pJmpRealAddressMemory, (void*)Asm_GetAddressQwordValueAsmFunctionStart(), Asm_GetAddressQwordValueLenAsmFunction());
						*(ULONG_PTR*)(pJmpRealAddressMemory + 0x2) = nRealAddress;
						KdPrint(("[Whisper] JmpRealAddress Offset: %I64X \n", pJmpRealAddressMemory - (pTempCodeBuffer + 5)));
						pTempCodeBuffer[0] = 0xE8;
						*(ULONG*)&(pTempCodeBuffer[1]) = (ULONG)(pJmpRealAddressMemory - (pTempCodeBuffer + 5));
						pTempCodeBuffer[5] = 0x90;
						pTempCodeBuffer[6] = 0x90;
					}
				}

				pTempCodeBuffer += nTempLen;
			}
			else {
				nCodeTempLen = 0;
			}
		}
		RtlCopyMemory(pJmpBackFunctionMemory + nCodeLen, (void *)pJmpBackFunctionStart, nJmpBackFunctionLen);
		*(ULONG_PTR *)(pJmpBackFunctionMemory + nCodeLen + 0x3) = pOriginalFunction + nCodeLen;

		ULONG_PTR nHookAgentFunctionLen = Asm_GetAgentAsmFunctionLen();
		ULONG_PTR pHookAgentFunctionStart = Asm_GetAgentAsmFunctionStart();

		unsigned char *pAgentHookMemory = (unsigned char *)ExAllocatePoolWithTag(NonPagedPool, nHookAgentFunctionLen, kHyperPlatformCommonPoolTag);
		if (pAgentHookMemory != NULL)
		{
			RtlCopyMemory((void *)pAgentHookMemory, (void *)pHookAgentFunctionStart, nHookAgentFunctionLen);
			*(ULONG_PTR *)(pAgentHookMemory + 0x1D) = pParamPointer;
			*(ULONG_PTR *)(pAgentHookMemory + 0x2B) = pDealFunction;
			*(ULONG_PTR *)(pAgentHookMemory + 0x5C) = (ULONG_PTR)pJmpBackFunctionMemory;
			
			ULONG_PTR pFunctionStart = Asm_GetJmpCodeStart();
			ULONG_PTR pFunctionLen = Asm_GetJmpCodeLen();
			RtlCopyMemory((void *)pFunction, (void *)pFunctionStart, pFunctionLen);
			*(ULONG_PTR *)(pFunction + 3) = (ULONG_PTR)pAgentHookMemory;
			bReturn = true;
		}
	}

	return bReturn;
}